package com.eshipenet.shipnet.edgeboxextbuilder

import org.apache.flink.api.common.state.MapStateDescriptor
import org.apache.flink.streaming.api.datastream.BroadcastStream
import org.apache.flink.streaming.api.scala._
import play.api.libs.json.{JsObject, JsValue, Json}
import com.github.nscala_time.time.Implicits._
import org.joda.time.DateTime

object Main {

  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment

    env.getConfig.registerTypeWithKryoSerializer(classOf[DateTime], classOf[DateTimeSerializer])

    val boxJsValueStream = createEdgeBoxStream(env)
    val boxActionStream = boxJsValueStream
      .map(jsValue => DebeziumParser(jsValue))
      .filter(_.isSuccess)
      .map(_.get)
      .filter(_.isDefined)
      .map(_.get)
      .map(EdgeBoxAction(_))
      .filter(_.isSuccess)
      .map(_.get)

    val modelJsValueStream = createEdgeBoxModelStream(env)
    val modelActionStream = modelJsValueStream
      .map(DebeziumParser(_))
      .filter(_.isSuccess)
      .map(_.get)
      .filter(_.isDefined)
      .map(_.get)
      .map(EdgeBoxModelAction(_))
      .filter(_.isSuccess)
      .map(_.get)

    val edgeBoxStream = mapEdgeBoxStream(boxActionStream)
      .filter(_.isDefined)
      .map(_.get)
      .keyBy(_.id)
      .mapWithState((action, state: Option[Map[String, Set[EdgeBox]]]) => {
        action match {
          case EdgeBoxInsert(id, after, timestamp) =>
            val newState = if (state.isEmpty) {
              Map[String, Set[EdgeBox]](after.modelId -> Set(after))
            } else {
              val sets = state.get.get(after.modelId) match {
                case Some(s) => s + after
                case _       => Set(after)
              }
              state.get.updated(after.modelId, sets)
            }
            val out = newState
            println("state", newState)
            (out, Some(newState))
          case EdgeBoxUpdate(id, before, after, timestamo) =>
            (Map[String, Set[EdgeBox]](), Some(Map[String, Set[EdgeBox]]()))
          case EdgeBoxDelete(id, before, timestamp) =>
            (Map[String, Set[EdgeBox]](), Some(Map[String, Set[EdgeBox]]()))
        }
      })
      .flatMap(d => d)

    val modelUpdateStream = mapModelChangeStream(edgeBoxStream, modelActionStream)
    modelUpdateStream.addSink(EdgeBoxExtActionSinkProvider.kafkaSink)

    val broadcastModelStream = mapBroadcastModelStream(modelActionStream)

    val boxUpdateStream = mapEdgeBoxAction(boxActionStream, broadcastModelStream)
//    boxUpdateStream.print()
    boxUpdateStream.addSink(EdgeBoxExtActionSinkProvider.kafkaSink)
    env.execute()
  }

  private def createEdgeBoxStream(env: StreamExecutionEnvironment): DataStream[JsObject] =
    EdgeBoxActionJsonStreamProvider.edgeBoxActionJsonStream(env)

  private def createEdgeBoxModelStream(env: StreamExecutionEnvironment): DataStream[JsObject] =
    EdgeBoxActionJsonStreamProvider.edgeBoxActionJsonStream(env)

  private def mapEdgeBoxStream(stream: DataStream[EdgeBoxAction]): DataStream[Option[EdgeBoxAction]] = {
    stream
      .keyBy(_.id)
      .mapWithState((action, state: Option[EdgeBoxAction]) => {
        if (state.isEmpty) {
          action match {
            case a: EdgeBoxInsert =>
              (Some(a), Some(a))
            case a: EdgeBoxUpdate =>
              (Some(a), Some(a))
            case a: EdgeBoxDelete =>
              (Some(a), Some(a))
          }
        } else {
          val currentTs = state.get.actionInfo.timestamp
          action match {
            case a @ EdgeBoxInsert(_, _, act) =>
              if (act.timestamp >= currentTs) {
                (Some(a), Some(a))
              } else {
                (None, state)
              }
            case a @ EdgeBoxUpdate(_, _, _, act) =>
              if (act.timestamp >= currentTs) {
                (Some(a), Some(a))
              } else {
                (None, state)
              }
            case a @ EdgeBoxDelete(_, _, act) =>
              if (act.timestamp >= currentTs) {
                (Some(a), Some(a))
              } else {
                (None, state)
              }
          }
        }
      })
  }

  private def mapModelChangeStream(stream: DataStream[(String, Set[EdgeBox])],
                                   modelStream: DataStream[EdgeBoxModelAction]) = {
    stream.connect(modelStream).keyBy(_._1, _.id).process(new EdgeBoxModelActionCoProcessFunction)
  }

  private lazy val modelStateDesc =
    new MapStateDescriptor[String, EdgeBoxModel]("EdgeBoxModel", classOf[String], classOf[EdgeBoxModel])

  private def mapBroadcastModelStream(stream: DataStream[EdgeBoxModelAction]) = {
    stream.broadcast(modelStateDesc)
  }

  private def mapEdgeBoxAction(stream: DataStream[EdgeBoxAction],
                               broadcastModelStream: BroadcastStream[EdgeBoxModelAction]) = {
    stream
      .connect(broadcastModelStream)
      .process(new EdgeBoxActionProcessFunction(modelStateDesc))
  }

}
